from plotting_functions import *

def makePlot(state_name, election_name):
    steps = 100000 # (this should accord to the length of your chain run)
    if state_name[0] == "n":
        state_abbr = "NC"
        bounds = [7,10]
        districts = 13
        pg_bound = 0.01
    elif state_name[0] == "t":
        state_abbr = "TX"
        bounds = [22,27]
        districts = 36
        pg_bound = 0.03
    elif state_name[0] == "u":
        state_abbr = "UT"
        bounds = [4,4]
        districts = 4
        pg_bound = 0.01

    datadir = "../outputs/" + state_abbr + "output/"

    newdir = "../plots/High_Res_Plots/"
    os.makedirs(os.path.dirname(newdir + "init.txt"), exist_ok=True)
    with open(newdir + "init.txt", "w") as f:
        f.write("Created Folder")

    max_steps = steps
    step_size = 10000 # this should accord to the dump_interval in the ensemble generating script

    ts = [x*step_size for x in range(1,int(max_steps/step_size)+1)]

    df = pd.DataFrame(columns = ['seats','mm','pg','vs','eg','ce'])

    for t in ts:
        tempdf = pd.read_csv(datadir + state_name + election_name +"_data"+str(t)+".csv", delimiter=',')
        df = pd.concat([df, tempdf], ignore_index=True)
        
    visualize(state_abbr+"-"+election_name, df, bounds, districts, pg_bound, 0.001, state_name, election_name, str(bounds), newdir)
    return

if __name__=="__main__":
    makePlot("texas", "SEN12")
    makePlot("northcarolina", "SEN16")
    makePlot("utah", "SEN16")